#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward Roualdes
#         [2012.12.13]
#         University of Kentucky

import utils
import sys
import os.path
from collections import defaultdict

def let2num(x):
    """change letters to numbers for appropriate STRUCTURE input."""
    # check for capital letters, since I don't know
    # if they will make it this far
    if x == 'a' or x == 'A':
        return str(1)
    elif x == 'c' or x == 'C':
        return str(2)
    elif x == 'g' or x == 'G':
        return str(3)
    elif x == 't' or x == 'T':
        return str(4)
    elif x == '-' or x == 'N' or x == 'n': 
        return str(-9)      
    else:
        print 'an unexpected letter!', x
        sys.exit()

def mkStruct(outdir, cleanSNP = None):
    """create structure file from inferred and aligned haplotypes"""

    print "Making STRUCTURE file..."
    bases = ['A', 'C', 'G', 'T', 'a', 'c', 'g', 't']
    allbases = bases+['-', 'n', 'N']

    # get aligned haplotypes data
    data = utils.files2dict(os.path.join(outdir, 'haplotypes'), '.*\\.aligned$')
    
    # find and store snps
    snp = {}                        # dict of {key = col names of snps.txt : val = {k = row_name : v = nucleotide}
    row_names = []                  # row names of snps.txt
    for d, indiv, locus in data:
        assert d, "no data in file %s.fna.aligned." % locus
        read_mat, indiv_names = utils.read_matrix(d)

        # search each row at column j for snp
        for j in xrange(len(read_mat[0])):
            col = [row[j] for row in read_mat if row[j] in allbases]
            if len([x for x in utils.unique(col) if x in bases]) > 1: # polymorphic position?
                # store locus and position of snp
                key = locus + ':' + str(j+1)
                snp[key] = dict(zip(indiv_names, 
                                    [utils.let2num(x) for x in col]))
                # track row_names
                row_names = utils.unique(row_names + indiv_names)

    # format snps
    txt = [['individual'] + row_names] # store data in columns
    for c in snp.keys():
        tmp = [c]
        for r in row_names:
            if r in snp[c].keys():
                tmp.append(snp[c][r])
            else:
                tmp.append(-9)
        txt.append(tmp)

    # clean SNPs
    if cleanSNP:
        print 'Cleaning SNPs...'
        for col in txt[1::]:
            if '-9' in col and len(utils.unique(col)) == 2:
                txt.remove(col)
    txt = zip(*txt) # transpose into rows of data; better for writing to file

    # write to file
    with open(os.path.join(outdir, 'snps.txt'), 'w') as f:
        for row in txt:
            for col in row:
                f.write('%s\t' % col)
            f.write('\n')
            
    return

if __name__ == "__main__":
    
    # outdir = '/Users/easy-e/Dropbox/uk/3.1/weisrock/ProjectNA/tiger_trim5_out'
    outdir = '/Users/easy-e/weisrock/nextallele/misc/test_cSNP'
    mkStructNew(outdir, True)
    print 'Done.'


